package geppetto.reordering;

import geppetto.cat.alignments.Alignment;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.phraseHMM.phraseExtraction.extractedphrase.ExtractedPhrasePair;
import geppetto.phraseTable.PhraseTable;
import geppetto.phraseTable.phrase.Phrase;
import gnu.trove.TIntDoubleHashMap;
import gnu.trove.TIntObjectHashMap;

import java.io.PrintStream;
import java.util.HashMap;

public class AlignmentWeightedMSDReorderingCalculator extends MSDReorderingCalculator {	
	
	protected REO_MODEL_TYPE modelType;
	protected TIntDoubleHashMap prevSmoothingFactors = new TIntDoubleHashMap();
	protected TIntDoubleHashMap nextSmoothingFactors = new TIntDoubleHashMap();
		
	public AlignmentWeightedMSDReorderingCalculator(REO_MODEL_TYPE modelType,AlignmentsSet set) {
		this(DEFAULTSMOOTHING,modelType, set);
	}
	
	public AlignmentWeightedMSDReorderingCalculator(double smoothing, REO_MODEL_TYPE modelType,AlignmentsSet set) {
		super(smoothing, modelType, set);
		this.modelType = modelType;
		this.set = set;
		setSmoothing(smoothing);
	}

	public AlignmentWeightedMSDReorderingCalculator(AlignmentsSet set) {
		this(REO_MODEL_TYPE.REO_MSD, set);
	}
	
	@Override
	public ReorderingInfo getReorderingInfo(ExtractedPhrasePair epp){
		 int sentenceNumber = epp.getSentenceNumber();
		 int startS = epp.getSourceStartPos();
		 int endS = epp.getSourceEndPos();
		 int startT = epp.getForeignStartPos();
		 int endT = epp.getForeignEndPos(); 
		 
		Alignment alignment = set.get(sentenceNumber);
		double connectedLeftTopP = getAlignmentProb(alignment, startS-1, startT-1);
		double connectedRightTopP = getAlignmentProb(alignment,endS+1, startT-1 );
		double connectedLeftTopN = getAlignmentProb(alignment,endS+1, endT+1 );
		double connectedRightTopN = getAlignmentProb(alignment,startS-1, endT+1 ); 
		
		double probLeftP = connectedLeftTopP * (1 - connectedRightTopP);
		double probRightP = connectedRightTopP * (1 - connectedLeftTopP);
		double probUnkP = connectedLeftTopP * connectedRightTopP + (1 - connectedLeftTopP) * (1 - connectedRightTopP);
		double probLeftN = connectedLeftTopN * (1 - connectedRightTopN);
		double probRightN = connectedRightTopN * (1 - connectedLeftTopN);
		double probUnkN = connectedLeftTopN * connectedRightTopN + (1 - connectedLeftTopN) * (1 - connectedRightTopN);

		ReorderingInfo info = new ReorderingInfo();
		info.addPrevWordOrientation(convertREO_POS(REO_POS.LEFT), probLeftP);
		info.addPrevWordOrientation(convertREO_POS(REO_POS.RIGHT), probRightP);
		info.addPrevWordOrientation(convertREO_POS(REO_POS.UNKNOWN), probUnkP);
		info.addNextWordOrientation(convertREO_POS(REO_POS.LEFT), probLeftN);
		info.addNextWordOrientation(convertREO_POS(REO_POS.RIGHT), probRightN);
		info.addNextWordOrientation(convertREO_POS(REO_POS.UNKNOWN), probUnkN);
		
		return info;
	}
	
	public int convertREO_POS(REO_POS orient){
		switch(orient){
		 	case LEFT: return MONO_CODE;
		 	case RIGHT: return SWAP_CODE;
		 	case DRIGHT: return DRIGHT_CODE;
		 	case DLEFT: return DLEFT_CODE;
		 	case UNKNOWN:
		 	switch(modelType){
		 	case REO_MONO: return NOMONO_CODE;
		 	case REO_MSD: return OTHER_CODE;
		 	case REO_MSLR: return DRIGHT_CODE;
		 	}
		 	break;
		 	} 
		return OTHER_CODE; // this should not happen
	}
	
	private double getAlignmentProb(Alignment alignment, int posS, int posT){
		if((-1 == posT && -1 == posS) || (alignment._foreignLen == posT && alignment._sourceLen == posS)){
			return 1;
		}
		else if(alignment.hasPosition(posS, posT)){
			return alignment._posteriors[posS][posT];
		}
		return 0;
	}

	
}
